from net import HRNet
from preprocess import get_data
import config
import logging
import torch
import torch.nn as nn

def test(model_path):
    logging.info("Start testing.")
    # get data
    testloader = get_data(config.data_path, train=False)

    net = HRNet()
    # net.init_weights(model_path)

    checkpoint = torch.load(model_path)
    net.load_state_dict(checkpoint['model_state_dict'])
    total = correct = 0
    for i, data in enumerate(testloader, 0):
        inputs, labels = data
        outputs = net(inputs)


        # get accuaracy of test dataset
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        logging.info('Accuracy of the network on the %d test images: %.3f %%' % (total,
                100.0 * correct / total))
    logging.info('Finished Training')
    
if __name__ == "__main__":
    logging.basicConfig(filename="log/test.log", filemode="w", format="%(asctime)s %(name)s:%(levelname)s:%(message)s", datefmt="%d-%M-%Y %H:%M:%S", level=logging.DEBUG)
    test("./models/saved_model.model")
